# models/enhanced_cnn_transformer.py
import torch
import torch.nn as nn
import torch.nn.functional as F


class EnhancedCNNTransformer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        input_channels = 4 if config.feature_type == "enhanced" else 2

        # CNN特征提取
        self.cnn = nn.Sequential(
            nn.Conv3d(input_channels, 32, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
            nn.BatchNorm3d(32),
            nn.GELU(),

            nn.Conv3d(32, 64, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
            nn.BatchNorm3d(64),
            nn.GELU(),

            nn.AdaptiveAvgPool3d((None, None, 100))  # [batch, 64, 3, 30, 100]
        )

        # Transformer时序建模
        self.transformer_encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=64 * 3 * 30,
                nhead=8,
                dim_feedforward=1024,
                dropout=config.dropout,
                batch_first=True
            ),
            num_layers=2
        )

        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(64 * 3 * 30, 512),
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(512, config.num_classes)
        )

    def forward(self, x):
        # CNN特征提取 [batch, C, 3, 30, T]
        x = self.cnn(x)

        # 重组为时序数据 [batch, T, D]
        batch_size = x.size(0)
        x = x.permute(0, 4, 1, 2, 3)  # [batch, T, C, 3, 30]
        x = x.reshape(batch_size, -1, 64 * 3 * 30)  # [batch, T, D]

        # Transformer处理
        x = self.transformer_encoder(x)

        # 全局平均分类
        x = x.mean(dim=1)
        return self.classifier(x)